import lightgbm as lgb
import pandas as pd
import numpy as np
import os
import joblib
import shap
from typing import Dict, List, Any, Tuple
from sklearn.metrics import (accuracy_score, precision_score, recall_score, 
                           f1_score, roc_auc_score, classification_report,
                           confusion_matrix)


class HardwareModel:
    """
    硬件故障预测模型
    """
    
    def __init__(self, config: Dict[str, Any]):
        self.config = config
        self.battery_model = None
        self.screen_model = None
        self.shap_explainer = None
    
    def train_battery_model(self, X_train: pd.DataFrame, y_train: pd.Series) -> None:
        """
        训练电池故障预测模型
        """
        params = {
            'objective': 'binary',
            'metric': 'auc',
            'boosting_type': 'gbdt',
            'learning_rate': self.config['training']['hardware']['learning_rate'],
            'max_depth': self.config['training']['hardware']['max_depth'],
            'num_leaves': self.config['training']['hardware']['num_leaves'],
            'random_state': self.config['training']['hardware']['random_state'],
            'n_estimators': self.config['training']['hardware']['num_iterations'],
            'verbose': -1
        }
        
        self.battery_model = lgb.LGBMClassifier(**params)
        self.battery_model.fit(X_train, y_train)
        
        # 创建SHAP解释器
        self.shap_explainer = shap.TreeExplainer(self.battery_model)
    
    def train_screen_model(self, X_train: pd.DataFrame, y_train: pd.Series) -> None:
        """
        训练屏幕故障预测模型
        """
        params = {
            'objective': 'regression',
            'metric': 'mse',
            'boosting_type': 'gbdt',
            'learning_rate': self.config['training']['hardware']['learning_rate'],
            'max_depth': self.config['training']['hardware']['max_depth'],
            'num_leaves': self.config['training']['hardware']['num_leaves'],
            'random_state': self.config['training']['hardware']['random_state'],
            'n_estimators': self.config['training']['hardware']['num_iterations'],
            'verbose': -1
        }
        
        self.screen_model = lgb.LGBMRegressor(**params)
        self.screen_model.fit(X_train, y_train)
    
    def evaluate_battery_model(self, X_test: pd.DataFrame, y_test: pd.Series) -> Dict[str, float]:
        """
        评估电池故障预测模型
        """
        if self.battery_model is None:
            raise ValueError("电池模型尚未训练")
        
        y_pred = self.battery_model.predict(X_test)
        y_pred_prob = self.battery_model.predict_proba(X_test)[:, 1]
        
        metrics = {
            'accuracy': accuracy_score(y_test, y_pred),
            'precision': precision_score(y_test, y_pred),
            'recall': recall_score(y_test, y_pred),
            'f1_score': f1_score(y_test, y_pred),
            'roc_auc': roc_auc_score(y_test, y_pred_prob)
        }
        
        print("电池故障预测模型评估:")
        for key, value in metrics.items():
            print(f"{key}: {value:.4f}")
        
        print("\n分类报告:")
        print(classification_report(y_test, y_pred))
        
        return metrics
    
    def evaluate_screen_model(self, X_test: pd.DataFrame, y_test: pd.Series) -> Dict[str, float]:
        """
        评估屏幕故障预测模型
        """
        if self.screen_model is None:
            raise ValueError("屏幕模型尚未训练")
        
        y_pred = self.screen_model.predict(X_test)
        
        # 计算回归指标
        mse = np.mean((y_pred - y_test) ** 2)
        rmse = np.sqrt(mse)
        mae = np.mean(np.abs(y_pred - y_test))
        r2 = np.corrcoef(y_pred, y_test)[0, 1] ** 2
        
        metrics = {
            'mse': mse,
            'rmse': rmse,
            'mae': mae,
            'r2': r2
        }
        
        print("屏幕故障预测模型评估:")
        for key, value in metrics.items():
            print(f"{key}: {value:.4f}")
        
        return metrics
    
    def get_shap_values(self, X: pd.DataFrame) -> np.ndarray:
        """
        获取SHAP值
        """
        if self.shap_explainer is None:
            raise ValueError("SHAP解释器尚未初始化")
        
        return self.shap_explainer.shap_values(X)
    
    def save_models(self, battery_path: str, screen_path: str) -> None:
        """
        保存模型
        """
        os.makedirs(os.path.dirname(battery_path), exist_ok=True)
        os.makedirs(os.path.dirname(screen_path), exist_ok=True)
        
        joblib.dump(self.battery_model, battery_path)
        joblib.dump(self.screen_model, screen_path)
        print(f"模型已保存至: {battery_path}, {screen_path}")
    
    def load_models(self, battery_path: str, screen_path: str) -> None:
        """
        加载模型
        """
        self.battery_model = joblib.load(battery_path)
        self.screen_model = joblib.load(screen_path)
        self.shap_explainer = shap.TreeExplainer(self.battery_model)
        print(f"模型已加载: {battery_path}, {screen_path}")


class SoftwareModel:
    """
    软件异常检测模型
    """
    
    def __init__(self, config: Dict[str, Any]):
        self.config = config
        self.model = None
        self.isolation_forest = None
    
    def train(self, X_train: pd.DataFrame, y_train: pd.Series) -> None:
        """
        训练软件异常检测模型
        """
        from sklearn.ensemble import IsolationForest
        
        # 训练LightGBM分类器
        params = {
            'objective': 'binary',
            'metric': 'auc',
            'boosting_type': 'gbdt',
            'learning_rate': self.config['training']['software']['learning_rate'],
            'max_depth': self.config['training']['software']['max_depth'],
            'num_leaves': self.config['training']['software']['num_leaves'],
            'random_state': self.config['training']['software']['random_state'],
            'n_estimators': self.config['training']['software']['num_iterations'],
            'verbose': -1
        }
        
        self.model = lgb.LGBMClassifier(**params)
        self.model.fit(X_train, y_train)
        
        # 训练Isolation Forest用于异常检测
        self.isolation_forest = IsolationForest(
            n_estimators=100,
            max_samples='auto',
            contamination=0.1,
            random_state=self.config['training']['software']['random_state']
        )
        self.isolation_forest.fit(X_train)
    
    def evaluate(self, X_test: pd.DataFrame, y_test: pd.Series) -> Dict[str, float]:
        """
        评估软件异常检测模型
        """
        if self.model is None:
            raise ValueError("模型尚未训练")
        
        y_pred = self.model.predict(X_test)
        y_pred_prob = self.model.predict_proba(X_test)[:, 1]
        
        metrics = {
            'accuracy': accuracy_score(y_test, y_pred),
            'precision': precision_score(y_test, y_pred),
            'recall': recall_score(y_test, y_pred),
            'f1_score': f1_score(y_test, y_pred),
            'roc_auc': roc_auc_score(y_test, y_pred_prob)
        }
        
        print("软件异常检测模型评估:")
        for key, value in metrics.items():
            print(f"{key}: {value:.4f}")
        
        print("\n分类报告:")
        print(classification_report(y_test, y_pred))
        
        return metrics
    
    def get_anomaly_score(self, X: pd.DataFrame) -> np.ndarray:
        """
        获取异常分数（结合LightGBM概率和Isolation Forest分数）
        """
        if self.model is None or self.isolation_forest is None:
            raise ValueError("模型尚未训练")
        
        # LightGBM的异常概率
        lgb_prob = self.model.predict_proba(X)[:, 1]
        
        # Isolation Forest的异常分数（转换为正数，越高越异常）
        if_score = -self.isolation_forest.score_samples(X)
        
        # 标准化Isolation Forest分数到[0,1]
        if_score = (if_score - if_score.min()) / (if_score.max() - if_score.min() + 1e-10)
        
        # 加权融合
        anomaly_score = 0.7 * lgb_prob + 0.3 * if_score
        
        return anomaly_score
    
    def save(self, model_path: str) -> None:
        """
        保存模型
        """
        os.makedirs(os.path.dirname(model_path), exist_ok=True)
        model_dict = {
            'lgb_model': self.model,
            'isolation_forest': self.isolation_forest
        }
        joblib.dump(model_dict, model_path)
        print(f"模型已保存至: {model_path}")
    
    def load(self, model_path: str) -> None:
        """
        加载模型
        """
        model_dict = joblib.load(model_path)
        self.model = model_dict['lgb_model']
        self.isolation_forest = model_dict['isolation_forest']
        print(f"模型已加载: {model_path}")


class TicketModel:
    """
    工单分类模型
    """
    
    def __init__(self, config: Dict[str, Any]):
        self.config = config
        self.model = None
        self.label_encoder = None
    
    def train(self, X_train: pd.DataFrame, y_train: pd.Series) -> None:
        """
        训练工单分类模型
        """
        from sklearn.preprocessing import LabelEncoder
        
        # 标签编码
        self.label_encoder = LabelEncoder()
        y_train_encoded = self.label_encoder.fit_transform(y_train)
        
        params = {
            'objective': 'multiclass',
            'metric': 'multi_error',
            'boosting_type': 'gbdt',
            'learning_rate': self.config['training']['ticket']['learning_rate'],
            'max_depth': self.config['training']['ticket']['max_depth'],
            'num_leaves': self.config['training']['ticket']['num_leaves'],
            'random_state': self.config['training']['ticket']['random_state'],
            'n_estimators': self.config['training']['ticket']['num_iterations'],
            'verbose': -1,
            'num_class': len(self.label_encoder.classes_)
        }
        
        self.model = lgb.LGBMClassifier(**params)
        self.model.fit(X_train, y_train_encoded)
    
    def evaluate(self, X_test: pd.DataFrame, y_test: pd.Series) -> Dict[str, float]:
        """
        评估工单分类模型
        """
        if self.model is None or self.label_encoder is None:
            raise ValueError("模型尚未训练")
        
        y_test_encoded = self.label_encoder.transform(y_test)
        y_pred_encoded = self.model.predict(X_test)
        
        metrics = {
            'accuracy': accuracy_score(y_test_encoded, y_pred_encoded),
            'precision_macro': precision_score(y_test_encoded, y_pred_encoded, average='macro'),
            'recall_macro': recall_score(y_test_encoded, y_pred_encoded, average='macro'),
            'f1_macro': f1_score(y_test_encoded, y_pred_encoded, average='macro')
        }
        
        print("工单分类模型评估:")
        for key, value in metrics.items():
            print(f"{key}: {value:.4f}")
        
        print("\n分类报告:")
        print(classification_report(y_test_encoded, y_pred_encoded, 
                                   target_names=self.label_encoder.classes_))
        
        return metrics
    
    def predict_proba(self, X: pd.DataFrame) -> pd.DataFrame:
        """
        预测概率
        """
        if self.model is None or self.label_encoder is None:
            raise ValueError("模型尚未训练")
        
        prob = self.model.predict_proba(X)
        prob_df = pd.DataFrame(prob, columns=self.label_encoder.classes_)
        return prob_df
    
    def save(self, model_path: str) -> None:
        """
        保存模型
        """
        os.makedirs(os.path.dirname(model_path), exist_ok=True)
        model_dict = {
            'model': self.model,
            'label_encoder': self.label_encoder
        }
        joblib.dump(model_dict, model_path)
        print(f"模型已保存至: {model_path}")
    
    def load(self, model_path: str) -> None:
        """
        加载模型
        """
        model_dict = joblib.load(model_path)
        self.model = model_dict['model']
        self.label_encoder = model_dict['label_encoder']
        print(f"模型已加载: {model_path}")


def train_hardware_model(config: Dict[str, Any], X_train: pd.DataFrame, y_train_battery: pd.Series, 
                        y_train_screen: pd.Series, X_test: pd.DataFrame, y_test_battery: pd.Series, 
                        y_test_screen: pd.Series) -> HardwareModel:
    """
    训练硬件模型的完整流程
    """
    model = HardwareModel(config)
    
    print("开始训练电池故障预测模型...")
    model.train_battery_model(X_train, y_train_battery)
    
    print("\n开始训练屏幕故障预测模型...")
    model.train_screen_model(X_train, y_train_screen)
    
    print("\n开始评估模型...")
    battery_metrics = model.evaluate_battery_model(X_test, y_test_battery)
    screen_metrics = model.evaluate_screen_model(X_test, y_test_screen)
    
    # 保存模型
    model.save_models(
        config['models']['hardware_model_path'].replace('.txt', '_battery.txt'),
        config['models']['hardware_model_path'].replace('.txt', '_screen.txt')
    )
    
    return model


def train_software_model(config: Dict[str, Any], X_train: pd.DataFrame, y_train: pd.Series, 
                         X_test: pd.DataFrame, y_test: pd.Series) -> SoftwareModel:
    """
    训练软件模型的完整流程
    """
    model = SoftwareModel(config)
    
    print("开始训练软件异常检测模型...")
    model.train(X_train, y_train)
    
    print("\n开始评估模型...")
    metrics = model.evaluate(X_test, y_test)
    
    # 保存模型
    model.save(config['models']['software_model_path'])
    
    return model


def train_ticket_model(config: Dict[str, Any], X_train: pd.DataFrame, y_train: pd.Series, 
                       X_test: pd.DataFrame, y_test: pd.Series) -> TicketModel:
    """
    训练工单分类模型的完整流程
    """
    model = TicketModel(config)
    
    print("开始训练工单分类模型...")
    model.train(X_train, y_train)
    
    print("\n开始评估模型...")
    metrics = model.evaluate(X_test, y_test)
    
    # 保存模型
    model.save(config['models']['ticket_model_path'])
    
    return model
